`timescale 1ns / 1ps

`define SLICE(BIT, IDX) (``BIT``) * (``IDX``) + (``BIT``) - 1 : (``BIT``) * (``IDX``)

// matmul (N_INH * N_INW) X (N_INW * N_OUT) @ 1 cycle
// N_INH * N_INW * N_OUT MACs @ 1 cycle
// N_INH * N_INW * N_OUT / 2 DSPs
module pwc2
#(
    parameter N_INH     = 1,    // N_IN  HEIGHT, must be 1, 2, 4, 8, ...
    parameter N_INW     = 16,   // N_IN  WIDTH
    parameter N_OUT     = 16,   // N_OUT WIDTH
    parameter N_ICH     = 384,
    parameter N_OCH     = 512,
    parameter BIT_I     = 8,
    parameter BIT_W     = 5,
    parameter BIT_O     = 16,
    parameter PE_TYPE   = "DSP",    // "DSP" or "LUT"
    parameter ROM_FILE  = "L9_PW.mem",
    parameter ROM_TYPE  = "block",
    parameter ROM_LATENCY = 2,
    parameter RAM_TYPE  = "distributed",
    parameter RAM_LATENCY = 1
)
(
    input   clk,
    input   rst,
    
    input   i_vld,
    output  i_rdy,
    input   [      N_INW*BIT_I-1 : 0]    i_data,
    
    output  o_vld,
    input   o_rdy,
    output  [N_INH*N_OUT*BIT_O-1 : 0]    o_data
);

localparam DATAPATH_LATENCY = (PE_TYPE == "LUT") ? 2 + N_INW : 3 + N_INW;
localparam W_IN  = N_INH * N_INW * BIT_I;
localparam W_OUT = N_OUT * BIT_O;
localparam W_WT  = N_OUT * N_INW * BIT_W;
localparam D_WT  = N_OCH * N_ICH / (N_OUT * N_INW);

wire  ii_vld;
wire  ii_last;
wire  [W_IN-1 : 0]  ii_data;
wire  wt_vld;
wire  [W_WT-1 : 0]  wt_data;

wire  [N_INW * BIT_I -1 : 0] ii_shft [N_INH -1 : 0];

wire in_rdy;
wire in_vld;
wire enable;

wire dy_vld;
wire dy_last;
wire dy_vld2;
wire ce_acc;
reg  ld_acc;    // load: 0; accumulate: 1

assign in_vld = ii_vld & wt_vld;
assign in_rdy = in_vld & enable;
assign enable = ~o_vld | o_rdy;
assign ce_acc = enable & dy_vld2;

pwc2_buffer #(
    .N_INH       ( N_INH        ),
    .N_INW       ( N_INW        ),
    .N_OUT       ( N_OUT        ),
    .N_ICH       ( N_ICH        ),
    .N_OCH       ( N_OCH        ),
    .BIT         ( BIT_I        ),
    .RAM_TYPE    ( RAM_TYPE     ),
    .RAM_LATENCY ( RAM_LATENCY  )
) inst_inbuf
(
    .clk                     ( clk      ),
    .rst                     ( rst      ),

    .i_vld                   ( i_vld    ),
    .i_data                  ( i_data   ),
    .i_rdy                   ( i_rdy    ),

    .o_rdy                   ( in_rdy   ),
    .o_vld                   ( ii_vld   ),
    .o_last                  ( ii_last  ),
    .o_data                  ( ii_data  )
);

zq_rom_loader #(
    .WIDTH          (W_WT),
    .DEPTH          (D_WT),
    .ROM_FILE       (ROM_FILE),
    .ROM_TYPE       (ROM_TYPE),
    .ROM_LATENCY    (ROM_LATENCY)
) weight_loader
(
    .clk        (clk),
    .rst        (rst),

    .o_vld      (wt_vld),
    .o_rdy      (in_rdy),
    .o_data     (wt_data)
);

shift_reg_rst #(
    .N          (DATAPATH_LATENCY - 2),
    .WIDTH      (2)
) inst_ctrl_line
(
    .clk        (clk),
    .rst        (rst),
    .clken      (enable),
    .i_data     ({in_vld, ii_last}),
    .o_data     ({dy_vld, dy_last})
);

always @(posedge clk)
begin
    if (rst)
        ld_acc <= 1'b0;
    else
    begin
        if (enable)
        begin
            if (dy_last)
                ld_acc <= 1'b0;
            else
                ld_acc <= ld_acc | dy_vld;
        end
    end
end

shift_reg_rst #(
    .N          (1),
    .WIDTH      (1)
) inst_pce_line
(
    .clk        (clk),
    .rst        (rst),
    .clken      (enable),
    .i_data     (dy_vld),
    .o_data     (dy_vld2)
);

shift_reg_rst #(
    .N          (2),
    .WIDTH      (1)
) inst_ovld_line
(
    .clk        (clk),
    .rst        (rst),
    .clken      (enable),
    .i_data     (dy_last),
    .o_data     (o_vld)
);

genvar i, j, k;
generate
    // data delay line
    for (i = 0; i < N_INH; i = i + 1)
    begin
        for (j = 0; j < N_INW; j = j + 1)
        begin
            shift_reg #(
                .N          ( j     ),
                .WIDTH      ( BIT_I )
            ) delay_line_data
            (
                .clk        ( clk    ),
                .clken      ( enable ),
                .i_data     ( ii_data   [ `SLICE(BIT_I, i * N_INW + j) ] ),
                .o_data     ( ii_shft[i][ `SLICE(BIT_I, j) ]             )
            );
        end
    end
    // DSP array
    if (PE_TYPE == "DSP") begin: DSP
        // weight delay line for DSP 2 in 1
        wire  [N_INW * 2 * BIT_W -1 : 0] wt_shft [N_OUT/2 -1 : 0];

        for (i = 0; i < N_OUT / 2; i = i + 1)
        begin
            for (j = 0; j < N_INW; j = j + 1)
            begin
                shift_reg #(
                    .N          ( j         ),
                    .WIDTH      ( 2 * BIT_W )
                ) delay_line_weight
                (
                    .clk        ( clk    ),
                    .clken      ( enable ),
                    .i_data     ( wt_data   [ `SLICE(2 * BIT_W, i * N_INW + j) ] ),
                    .o_data     ( wt_shft[i][ `SLICE(2 * BIT_W, j) ]             )
                );
            end
        end

        for (i = 0; i < N_INH; i = i + 1)
        begin: INH
            for (j = 0; j < N_OUT / 2; j = j + 1)
            begin: OUT
                pwc2_dspchain #(
                    .N_PE   ( N_INW  ),
                    .BIT_I  ( BIT_I  ),
                    .BIT_W  ( BIT_W  ),
                    .BIT_O  ( BIT_O  ),
                    .LSHIFT ( 20     )
                ) inst_dspchain
                (
                    .clk                     ( clk        ),
                    .ce                      ( enable     ),
                    .ce_acc                  ( ce_acc     ),
                    .ld_acc                  ( ld_acc     ),
                    .i_data                  ( ii_shft[i] ),
                    .i_weight                ( wt_shft[j]  ),
                    .o_data                  (  o_data[ `SLICE(2 * BIT_O, i * N_OUT/2 + j) ] )
                );
            end
        end
    end
    else if (PE_TYPE == "LUT") begin: LUT
        // weight delay line
        wire  [N_INW * BIT_W -1 : 0] wt_shft [N_OUT -1 : 0];

        for (i = 0; i < N_OUT; i = i + 1)
        begin
            for (j = 0; j < N_INW; j = j + 1)
            begin
                shift_reg #(
                    .N          ( j      ),
                    .WIDTH      ( BIT_W  )
                ) delay_line_weight
                (
                    .clk        ( clk    ),
                    .clken      ( enable ),
                    .i_data     ( wt_data   [ `SLICE(BIT_W, i * N_INW + j) ] ),
                    .o_data     ( wt_shft[i][ `SLICE(BIT_W, j) ]             )
                );
            end
        end

        for (i = 0; i < N_INH; i = i + 1)
        begin: INH
            for (j = 0; j < N_OUT; j = j + 1)
            begin: OUT
                pwc2_pechain #(
                    .N_PE  ( N_INW  ),
                    .BIT_I ( BIT_I  ),
                    .BIT_W ( BIT_W  ),
                    .BIT_O ( BIT_O  )
                ) inst_pechain (
                    .clk                     ( clk        ),
                    .ce                      ( enable     ),
                    .ce_acc                  ( ce_acc     ),
                    .ld_acc                  ( ld_acc     ),
                    .i_data                  ( ii_shft[i] ),
                    .i_wigt                  ( wt_shft[j] ),
                    .o_data                  (  o_data[ `SLICE(BIT_O, i * N_OUT + j) ] )
                );
            end
        end
    end
endgenerate

endmodule
